import os
import numpy as np
import matplotlib.pyplot as plt
from scipy.interpolate import UnivariateSpline
from scipy.optimize import fsolve
from scipy.signal import argrelextrema
"""
from multiprocessing import cpu_count
from concurrent.futures import ProcessPoolExecutor
"""

save_path = os.getcwd()

omega_0 = 3
last_time = 30
gamma_precision = 1e-3

# 设置注释文本的样式与箭头样式
BBOX = dict(boxstyle="round", fc="0.8")
ARROW = dict(
    arrowstyle="->",
    connectionstyle="angle,angleA=0,angleB=90,rad=10")

def Simulation(position, velocity, delta_t, gamma, sim_method):
    acceleration = -np.square(omega_0) * position - gamma * velocity
    if sim_method == 'EulerCromer':
        velocity += acceleration * delta_t
        position += velocity * delta_t
    elif sim_method == 'Euler':
        position += velocity * delta_t
        velocity += acceleration * delta_t
    elif sim_method == 'EulerRichardson':
        position_half = position + velocity * delta_t / 2
        velocity_half = velocity + acceleration * delta_t / 2
        acceleration_half = -np.square(omega_0) * position_half - gamma * velocity_half
        position += velocity_half * delta_t
        velocity += acceleration_half * delta_t
    return position, velocity

def Calculation(sim_method, gamma):
    if sim_method == "Eular":
        precision_default = 300000
    else:
        precision_default = 30000
    moments = np.linspace(0, last_time, precision_default)
    delta_t = last_time / precision_default
    position = 1
    velocity = 0
    positions = [position]
    for moment in moments[:-1]:
        position, velocity = Simulation(position, velocity, delta_t, gamma, sim_method)
        positions.append(position)
    return moments, positions

def task1(sim_method):
    moments, positions = Calculation(sim_method, 0.5)
    moments = np.array(moments)
    positions = np.array(positions)
    result = []
    positions_greater_arg = np.array(argrelextrema(positions, np.greater))[0]
    positions_less_arg = np.array(argrelextrema(positions, np.less))[0]
    for i in range(2):
        func = UnivariateSpline(moments[positions_less_arg[i]:positions_greater_arg[i]], positions[positions_less_arg[i]:positions_greater_arg[i]])
        result.append(fsolve(func, 0)[0])
        plt.annotate(
            "$t=%fs$\n$x=%fm$"%(result[i], 0),
            xy=(result[i], 0),
            xytext=(20, (80 + 10 * i) * np.power(-1, i)),
            textcoords='offset points',
            bbox=BBOX,
            arrowprops=ARROW
        )
    plt.grid()
    plt.title(sim_method)
    plt.plot(moments, positions)
    plt.scatter(result, [0, 0])
    plt.savefig(os.path.join(save_path, "1_%s.png"%sim_method))
    plt.xlabel("$t/s$")
    plt.ylabel("$x/m$")
    plt.cla()

def task2(sim_method):
    gamma_1 = 0.5
    gamma_2 = 10
    gamma_3 = 0.5 * (gamma_1 + gamma_2)
    moments, positions = Calculation(sim_method, gamma_3)
    while gamma_2 - gamma_1 > gamma_precision:
        if np.min(positions) < 0:
            gamma_1 = gamma_3
        else:
            gamma_2 = gamma_3
        gamma_3 = 0.5 * (gamma_1 + gamma_2)
        moments, positions = Calculation(sim_method, gamma_3)
    plt.grid()
    plt.title(sim_method + ",$\gamma=%f$"%gamma_3)
    plt.plot(moments, positions)
    plt.savefig(os.path.join(save_path, "2_%s.png"%sim_method))
    plt.xlabel("$t/s$")
    plt.ylabel("$x/m$")
    plt.cla()
    

if __name__ == "__main__":
    """
    with ProcessPoolExecutor(max_workers=cpu_count()) as executor:
        for sim_method in ["Euler", "EulerCromer", "EulerRichardson"]:
            executor.submit(task1, sim_method)
            executor.submit(task2, sim_method)
    """
    for sim_method in ["Euler", "EulerCromer", "EulerRichardson"]:
        task1(sim_method)
        task2(sim_method)